import argparse
from functools import cached_property
import json
import logging
import os
from pathlib import Path
import tempfile
import polars as pl
from pydantic import BaseModel

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("CLNSIG")


class Clnsig(BaseModel):
    input_xlsx: Path
    output_tsv: Path
    taxdump_dir: Path

    @cached_property
    def outdir(self) -> Path:
        outdir = self.output_tsv.parent
        outdir.mkdir(parents=True, exist_ok=True)
        return outdir

    @cached_property
    def taxonkit_list_json(self) -> Path:
        return self.outdir / "taxonkit_list.json"

    @cached_property
    def taxonkit_lineage_tsv(self) -> Path:
        return self.outdir / "taxonkit_lineage.tsv"

    @cached_property
    def clnsig_records(self) -> pl.DataFrame:
        return (
            pl.read_excel(self.input_xlsx, columns=[1, 5, 6, 7])
            .rename({"TaxID": "taxid", "Blood": "blood", "Respiratory": "respiratory", "CerebrospinalFluid": "cerebrospinal_fluid"})
            .with_columns(pl.col("taxid").cast(pl.String).alias("taxid"))
        )

    @cached_property
    def taxid_records(self) -> pl.DataFrame:
        def get_sub_taxids(data: dict[str, dict]):
            taxids = []
            for taxid, sub_data in data.items():
                if sub_data:
                    taxids.extend(get_sub_taxids(sub_data))
                taxids.append(taxid)
            return taxids

        with open(self.taxonkit_list_json) as reader:
            records = pl.DataFrame(map(lambda x: {"taxid": x[0], "sub_taxid": [x[0]] + get_sub_taxids(x[1])}, json.load(reader).items()))
            return records.explode(["sub_taxid"]).unique()

    def taxonkit_list(self):
        with tempfile.NamedTemporaryFile(suffix=".txt") as temp_file:
            self.clnsig_records.select(["taxid"]).write_csv(temp_file.name, include_header=False)
            command = f"taxonkit list --data-dir {self.taxdump_dir} -J -j 16 -o {self.taxonkit_list_json} {temp_file.name}"
            logger.info("Run command: %s", command)
            if os.system(command) > 0:
                raise RuntimeError(f"Run 'taxonkit list' Fail: {command}")

    def taxonkit_lineage(self):
        with tempfile.NamedTemporaryFile(suffix=".txt") as temp_file:
            self.taxid_records.select(["sub_taxid"]).write_csv(temp_file.name, include_header=False)
            command = f"taxonkit filter --data-dir {self.taxdump_dir} -j 16 -E species -L species {temp_file.name} | \
                        taxonkit lineage --data-dir {self.taxdump_dir} -r  -j 16 -o {self.taxonkit_lineage_tsv}"
            logger.info("Run command: %s", command)
            if os.system(command) > 0:
                raise RuntimeError(f"Run 'taxonkit lineage' Fail: {command}")

    def write_out(self):
        records = pl.read_csv(self.taxonkit_lineage_tsv, separator="\t", new_columns=["taxid", "name", "rank"], has_header=False).cast(pl.String)
        records = records.with_columns(pl.col("name").str.split(";").list.last().alias("name"))
        records = records.join(self.taxid_records.join(self.clnsig_records, on="taxid", how="left").drop(["taxid"]), left_on="taxid", right_on="sub_taxid")
        records.write_csv(self.output_tsv, separator="\t")

    def run(self):
        if not self.taxonkit_list_json.exists():
            self.taxonkit_list()
        if not self.taxonkit_lineage_tsv.exists():
            self.taxonkit_lineage()
        self.write_out()


pathogen_dir = Path("/mnt/nas/database/pathogen_2/pathogen_genomics/")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Pathogen Genome")
    parser.add_argument("-t", "--taxdump_dir", type=Path, default=pathogen_dir / "taxdump", help="taxdump directory")
    parser.add_argument("-i", "--input_xlsx", type=Path, help="input excel file")
    parser.add_argument("-o", "--output_tsv", type=Path, help="output TSV file")
    # parser.add_argument("-n", "--nt_fasta", type=Path, default=Path("/mnt/nas/database/NT/nt.fasta"), help="NT file")
    # parser.add_argument("-w", "--workdir", type=Path, default=pathogen_dir, help="output directory")
    parser.set_defaults(func=lambda args: Clnsig.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
